package com.xiashitech.agent.instrumentation.biz.userinfo.springsecurity.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.constants.SpecialConfigString;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.biz.userinfo.springsecurity.AuthenticationData;
import com.xiashitech.agent.instrumentation.biz.userinfo.springsecurity.XSAuthenticationSingletons;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSAuthenticationCallDepth;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.agent.utils.XSUtil;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.named;

public class XSAuthenticationInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("org.springframework.security.core.Authentication");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return hasSuperType(named("org.springframework.security.core.Authentication"));
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("getPrincipal")),
                this.getClass().getName() + "$AuthenticationAdvice");
    }


    @SuppressWarnings("unused")
    public static class AuthenticationAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("otelBaggageScope") Scope baggageScope,
                @Advice.Local("authenticationData") AuthenticationData authenticationData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSAuthenticationInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                callDepth = CallDepth.forClass(XSAuthenticationCallDepth.class);
                if (callDepth.getAndIncrement() > 0) {
                    return;
                }

                String userId = Java8Bridge.currentForBaggage().getEntryValue("xs.userId");
                if (userId != null && !"".equals(userId)){
                    return;
                }
                String userIdPath = XSUtil.getAgentConfigBySystemId(AgentConfigDTO.getConfig().getUseridpath());
                if (XSUtil.isEmpty(userIdPath) || !userIdPath.startsWith(SpecialConfigString.BIZ_USER_SS.getConfiString())){
                    return;
                }
                Context parentContext = currentContext();
                authenticationData = AuthenticationData.builder().build();
                parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                context = XSAuthenticationSingletons.instrumenter().start(parentContext, authenticationData);
                context = BaggageUtil.syncBaggage().build().storeInContext(context);
                scope = context.makeCurrent();
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("AuthenticationAdvice b");
        }


        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(@Advice.Return Object returnObj,
                                  @Advice.Local("otelCallDepth") CallDepth callDepth,
                                  @Advice.Local("otelContext") Context context,
                                  @Advice.Local("otelScope") Scope scope,
                                  @Advice.Local("user") UserRequest userRequest,
                                  @Advice.Local("authenticationData") AuthenticationData authenticationData) {
            DurationPrinter durationPrinter = new DurationPrinter("XSAuthenticationInstrumentation");
            try {
                if(CacheStatus.cannotRun()) {
                    return;
                }
                if (callDepth.decrementAndGet() > 0) {
                    return;
                }
                if (scope == null) {
                    return;
                }
                String userId = Java8Bridge.currentForBaggage().getEntryValue("xs.userId");
                if (userId != null && !"".equals(userId)){
                    return;
                }
                String userIdPath = XSUtil.getAgentConfigBySystemId(AgentConfigDTO.getConfig().getUseridpath());
                if (XSUtil.isEmpty(userIdPath) || !userIdPath.startsWith(SpecialConfigString.BIZ_USER_SS.getConfiString())){
                    return;
                }
                authenticationData.setPrincipal(returnObj);
                scope.close();
                XSAuthenticationSingletons.instrumenter().end(context, authenticationData, null, null);
                // 理论上不需要这里再次赋值，但实际项目实施中发现有可能需要这个补丁。
                if (authenticationData.getUserId() != null){
                    Java8Bridge.currentForTraceContextVars()
                            .set("xs.userId", authenticationData.getUserId());
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("AuthenticationAdvice e");
        }

    }

}
